# -*- coding: utf-8 -*-

# __date:       2021/6/10
# __author:     YangChao
# __function:

import re
import numpy as np
from .visa_inst import VisaInst


def str2num(Mdata):
    try:
        Pydata = eval(Mdata)
    except:
        if re.findall(';', Mdata):
            new_data = Mdata.split(';')  # MATLAB中的分号表示一行数组
            Pydata = []
            for var in new_data:
                if re.findall(',', var):
                    order = var.split(',')
                else:
                    order = var.split()
                Pydata.append([eval(s) for s in order])
        else:
            if re.findall(',', Mdata):
                new_data = Mdata.split(',')
                Pydata = []
                for var in new_data:
                    Pydata.append(eval(var))
            else:
                new_data = Mdata.split()
                Pydata = []
                for var in new_data:
                    Pydata.append(eval(var))

    Pydata = np.array(Pydata)
    return Pydata


class SA4051C(VisaInst):

    def __init__(self, visa_address: str = 'TCPIP0::10.10.11.198::5025::SOCKET'):
        super().__init__(visa_address)
        return

    def rst(self):
        self.inst_obj.write('*RST\n')

    def set_external_trig(self, channel: int):
        order = 'TRIG:SOUR EXT%d' % channel
        self.inst_obj.write(order)

    def set_RBW(self, bw, f='KHZ'):
        # 设置RBW
        order = 'BAND %g %s' % (bw, f)
        self.inst_obj.write(order)

    def set_VBW(self, bw, f='KHZ'):
        # 设置VBW
        order = 'BAND:VID %g %s' % (bw, f)
        self.inst_obj.write(order)

    def set_trace_type(self, num: int, typ: str):
        order = 'TRAC%g:TYPE %s' % (num, typ)
        self.inst_obj.write(order)

    def set_span(self, span):
        order = 'FREQ:SPAN %g' % span
        self.inst_obj.write(order)

    def set_band_span(self, num, span):
        order = 'CALC:MARK%g:FUNC:BAND:SPAN %f' % (num, span)
        self.inst_obj.write(order)

    def set_center(self, cent):
        # 设置中心频率
        order = 'FREQ:CENT %g' % cent
        self.inst_obj.write(order)

    def set_marker(self, num: int, mode: str = 'Pos', position: float = None, func: str = None):
        # set marker mode 设置标记模式
        order_mode = 'CALC:MARK%d:MODE %s' % (num, mode)
        self.inst_obj.write(order_mode)
        # set marker position
        if position is not None:
            # 设置标记x轴位置
            order_pos = ':CALC:MARK%d:X %g' % (num, position)
            self.inst_obj.write(order_pos)
        # set marker function
        if func is not None:
            # 设置标记功能
            order_func = 'CALC:MARK%g:FUNC %s' % (num, func)
            self.inst_obj.write(order_func)

    def set_sweep_time(self, time, f: str = 's'):
        # 设置扫描时间
        order = 'SWE:TIME %g %s' % (time, f)
        self.inst_obj.write(order)

    def set_average_num(self, repeats: int):
        min_value = 1
        max_value = 10000
        if min_value <= repeats <= max_value:
            # 设置平均次数
            order = 'AVERage:COUNt %d' % repeats
            self.inst_obj.write(order)
        else:
            raise ValueError('repeats {} out of range {} ~ {}'
                             .format(repeats, min_value, max_value))

    # 获取marker幅值
    def get_marker(self, num):
        order = str(':CALC:MARK%d:Y?' % num)
        maker_amp = self.inst_obj.query(order)
        MarkerAmp = str2num(maker_amp)
        return MarkerAmp

    # 关闭所有Marker
    def marker_off(self):
        self.inst_obj.write(':CALC:MARK:AOFF\n')

    # 比较Marker的相对大小
    def delta(self, marker1, marker2):
        self.inst_obj.write('CALC:MARK%d:MODE DELT\n' % marker1)

    # 自动寻峰
    def peak(self):
        data = []
        self.inst_obj.write('CALC:MARK1:MAX\n')
        data.append(str2num(self.inst_obj.query('CALC:MARK1:Y?')))
        self.inst_obj.write('CALC:MARK2:MAX\n')
        self.inst_obj.write('CALC:MARK2:MAX:NEXT\n')
        data.append(str2num(self.inst_obj.query('CALC:MARK2:Y?')))
        self.inst_obj.write('CALC:MARK3:MAX\n')
        self.inst_obj.write('CALC:MARK3:MAX:NEXT\n')
        self.inst_obj.write('CALC:MARK3:MAX:NEXT\n')
        data.append(str2num(self.inst_obj.query('CALC:MARK3:Y?')))
        data = np.sort(data)
        self.set_marker(1, data[0])
        self.set_marker(2, data[1])
        self.set_marker(3, data[2])
        return np.array(data)

    # 中心频率的幅度
    def get_center_amp(self, aveTimes=1):
        if aveTimes == 1:
            _, ampList = self.get()
            # print(ampList)
            amp = ampList[int((len(ampList) + 1) / 2)]
        elif aveTimes > 1:
            ampforAve = np.zeros((aveTimes, 1))
            for i in range(aveTimes):
                _, ampList = self.get()
                amp = ampList[int((len(ampList) + 1) / 2)]
                ampforAve[i] = amp
            amp = np.mean(ampforAve)
        else:
            return
        return amp

    # 设置一个中心频率，得到幅度，一次，不平均
    def get_amp_by_freq(self, freq):
        self.set_center(freq)
        _, ampList = self.get()
        amp = ampList[int((len(ampList) + 1) / 2)]
        return amp

    def get_channel_power(self):
        order = 'CALCulate:DATA?'
        data_str = self.inst_obj.query(order)
        data = str2num(data_str)
        power = data[0]
        psd = data[1]
        return power, psd

    def get(self):
        # 具体的数据返回格式要通过调试仪器得到
        # 这里先判断数据是一维还是二维
        order = 'FETCh:SAN1?'
        # order = 'READ:SAN1?\n'
        data_str = self.inst_obj.query(order)
        data = str2num(data_str)
        # print(data)
        if len(data.shape) > 1:  # 多维数组
            freq = np.array([i for k in [v[0: len(v): 2] for v in data.T] for i in k]).conj().transpose()
            amp = np.array([i for k in [v[1: len(v): 2] for v in data.T] for i in k]).conj().transpose()
        else:
            freq = data[0:len(data):2].conj().transpose()
            amp = data[1:len(data):2].conj().transpose()
        return freq, amp

    # 设置开始频率
    def set_start_fre(self, fre):
        order = 'FREQ:STAR %g' % fre
        self.inst_obj.write(order)

    # 设置结束频率
    def set_end_fre(self, fre):
        order = 'FREQ:STOP %g' % fre
        self.inst_obj.write(order)

    def set_peak(self):
        self.inst_obj.write('CALCulate:MARK:PEAK:SEARch:MODE MAXimum')
        self.inst_obj.write('CALC:MARK:MAX\n')

    # 获得峰值处x（频率）的值
    def get_x_peak(self):
        data = str2num(self.inst_obj.query('CALC:MARK1:X?'))
        print(data)
        return data
